#include "myplatform.h"
#include "fl.h"
#include "diskprotect.h"
#define DISKFT_POOL_TAG 'ksid'
#include "dispatch.h"
#include "pnp.h"
#include "write_cache.h"
#pragma code_seg("PAGE") // start PAGE section

VOID diskftUnload(
    IN  PDRIVER_OBJECT  DriverObject
)
{

    g_bShutDown=NSS_TRUE;
    if(g_bHasProtectDisk)
        KeWaitForSingleObject(&g_ShutDownEvent,Executive,KernelMode,FALSE,NULL);
    NSS_PRINT("++");
    NSS_PRINT("--");
    return;
}
//////////////////////////////////////////////////////////////////////////////
#pragma code_seg() // end PAGE section

int GetRegValue(PCWSTR RegPath,PCWSTR ValueName,PULONG Value)
{
    int ReturnValue = 0;
    NTSTATUS Status;
    OBJECT_ATTRIBUTES ObjectAttributes;
    HANDLE KeyHandle;
    KEY_VALUE_PARTIAL_INFORMATION *Information;
    ULONG InformationSize;
    UNICODE_STRING UnicodeRegPath;
    UNICODE_STRING UnicodeValueName;

    RtlInitUnicodeString(&UnicodeRegPath, RegPath);
    RtlInitUnicodeString(&UnicodeValueName, ValueName);

    InitializeObjectAttributes(&ObjectAttributes,
                               &UnicodeRegPath,
                               OBJ_CASE_INSENSITIVE, // Flags
                               NULL, // Root directory
                               NULL); // Security descriptor

    Status = ZwOpenKey(&KeyHandle,
                       KEY_ALL_ACCESS,
                       &ObjectAttributes);
    if (Status != STATUS_SUCCESS)
    {
        NSS_PRINT("open key faild\n");
        return 0;
    }

    InformationSize = sizeof(Information) + sizeof(ULONG);
    Information = (KEY_VALUE_PARTIAL_INFORMATION*)ExAllocatePoolWithTag(PagedPool, InformationSize,'ISmD'); // DmSI

    if (Information == NULL)
    {
        ZwClose(KeyHandle);
        NSS_PRINT("Information==NULL\n");
        return 0;
    }

    Status = ZwQueryValueKey(KeyHandle,
                             &UnicodeValueName,
                             KeyValuePartialInformation,
                             Information,
                             sizeof(Information),
                             &InformationSize);
    if (Status == STATUS_SUCCESS)
    {
        if (Information->Type == REG_DWORD && Information->DataLength == sizeof(ULONG))
        {
            RtlCopyMemory(Value, Information->Data, sizeof(ULONG));
            ReturnValue = 1;
        }
        else
            NSS_PRINT("Infomation type!=DWORD\n");
    }
    else
        NSS_PRINT("status!=success\n");
    ExFreePool(Information);
    ZwClose(KeyHandle);
    return ReturnValue;
}

#pragma code_seg("INIT")
extern "C"
NTSTATUS DriverEntry(	IN PDRIVER_OBJECT DriverObject,
                        IN PUNICODE_STRING RegistryPath)
{
    NTSTATUS status = STATUS_SUCCESS;
    UNICODE_STRING str;
    ULONG i;

    g_bShutDown=NSS_FALSE;
    g_DiskCount=0;
    g_bHasProtectDisk=NSS_FALSE;

    KeInitializeEvent(&g_ShutDownEvent,SynchronizationEvent,FALSE);
    Array_Init();

    for (i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; ++i)
    {
        DriverObject->MajorFunction[i] = diskftPassThrough;
    }

    DriverObject->DriverExtension->AddDevice = diskftAddDevice;
    DriverObject->DriverUnload = diskftUnload;
    DriverObject->MajorFunction[IRP_MJ_SCSI] = diskftscsiDispatch;
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = diskftDeviceIoControlDispatch;
    DriverObject->MajorFunction[IRP_MJ_POWER] = Wdm1Power;
    DriverObject->MajorFunction[IRP_MJ_PNP] = Wdm1Pnp;

    NSS_PRINT("DriverEntry completed the vesion is 2.0");
    return status;
}
#pragma code_seg()

